# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-2.14
import abc
from typing import Any
from langchain.tools import BaseTool
import logging
import sys
import os
import pandas as pd
import utils.config_util as utils
from langchain_openai import ChatOpenAI
from langchain_community.embeddings.openai import OpenAIEmbeddings
import sqlite3
import utils.configFinRAG as configFinRAG
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.instances import TOKENIZER, LLM

# 定义全局变量
g_example_question_list = []
g_example_sql_list = []
g_example_fa_list = []
g_example_info_list = []
g_example_token_list = []

# 配置日志
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger(__name__)
logger.addHandler(logging.StreamHandler(stream=sys.stdout))

def load_sql_examples():
    """
    加载 SQL 示例数据
    """
    global g_example_question_list, g_example_sql_list, g_example_fa_list, g_example_info_list, g_example_token_list
    try:
        if not g_example_question_list:
            sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
            for _, row in sql_examples_file.iterrows():
                g_example_question_list.append(row['问题'])
                g_example_sql_list.append(row['SQL'])
                g_example_info_list.append(row['资料'])
                g_example_fa_list.append(row['FA'])
                tokens = TOKENIZER(row['问题'])
                g_example_token_list.append(tokens['input_ids'])
    except FileNotFoundError:
        logger.error(f"未找到文件: {configFinRAG.sql_examples_path}")
    except Exception as e:
        logger.error(f"加载 SQL 示例数据时发生错误: {e}")

class FinSQLRAG(BaseTool, abc.ABC):
    name: str = "查询金融数据库"
    description: str = "当被问到金融查询相关的问题时，会去金融数据库检索结果"

    def __init__(self):
        super().__init__()
        # 初始化时加载 SQL 示例数据
        load_sql_examples()

    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
        # 用例中没有用到 arun 不予具体实现
        pass

    def _run(self, para) -> str:
        query = para
        try:
            result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list,
                                              g_example_token_list)
            logger.info(f"query sql is: {sql}")

            # 数据库文件路径
            db_path = r'C:\Users\langzi\Desktop\local_human_number\数字人\Fay-fay-agent-edition0830\agent\tools\data\博金杯比赛数据.db'
            # 使用 with 语句管理数据库连接
            with sqlite3.connect(db_path) as conn:
                cs = conn.cursor()
                success_flag, exc_result = query_db(sql, cs)

            answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list,
                                     g_example_fa_list, g_example_token_list)
            return answer
        except sqlite3.Error as e:
            logger.error(f"数据库查询错误: {e}")
            return 'FinSQLRAG 数据库查询异常！'
        except Exception as e:
            logger.error(f"Retriever Error: {e}")
            return 'FinSQLRAG 处理异常！'

if __name__ == "__main__":
    tool = FinSQLRAG()
    question = "请帮我计算，在20210715，中信行业分类划分的一级行业为消费者服务行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。"
    result = tool.run(question)
    print(result)
